package com.tuniu.agents.advisor;

import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.util.ChatResponseUtils;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Agent LLM Chat History Interceptor
 * This class manages the chat memory and conversation history for AI agents
 */
public class AgentChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemory>{

    /**
     * Key for storing agent ID in chat memory
     */
    public static final String CHAT_MEMORY_AGENT_ID_KEY = "chat_memory_agent_id";

    /**
     * Key for storing conversation round in chat memory
     */
    public static final String CHAT_MEMORY_AGENT_ROUND = "chat_memory_agent_round";

    private static final String DEFAULT_SYSTEM_TEXT_ADVISE = "";

//    private static final String DEFAULT_SYSTEM_TEXT_ADVISE = """
//
//			Use the conversation memory from the MEMORY section to provide accurate answers.
//
//			---------------------
//			MEMORY:
//			{{memory}}
//			---------------------
//
//			""";

    private final String systemTextAdvise;

    /**
     * Flag to control if incoming chat messages should be recorded
     */
    private final boolean chatMemoryInRecord;

    /**
     * Flag to control if outgoing chat messages should be recorded
     */
    private final boolean chatMemoryOutRecord;

    /**
     * Constructor with default settings
     * @param chatMemory The chat memory instance to use
     */
    public AgentChatMemoryAdvisor(ChatMemory chatMemory) {
        this(chatMemory, DEFAULT_SYSTEM_TEXT_ADVISE, true,true); // Default value is true
    }

    /**
     * Constructor with memory recording control
     * @param chatMemory The chat memory instance to use
     * @param chatMemoryInRecord Flag to control incoming message recording
     * @param chatMemoryOutRecord Flag to control outgoing message recording
     */
    public AgentChatMemoryAdvisor(ChatMemory chatMemory, boolean chatMemoryInRecord,boolean chatMemoryOutRecord) {
        this(chatMemory, DEFAULT_SYSTEM_TEXT_ADVISE, chatMemoryInRecord,chatMemoryOutRecord);
    }

    /**
     * Constructor with custom system text
     * @param chatMemory The chat memory instance to use
     * @param systemTextAdvise Custom system text advice
     */
    public AgentChatMemoryAdvisor(ChatMemory chatMemory, String systemTextAdvise) {
        this(chatMemory, systemTextAdvise, true,true); // Default value is true
    }

    /**
     * Full constructor with all parameters
     * @param chatMemory The chat memory instance to use
     * @param systemTextAdvise Custom system text advice
     * @param chatMemoryInRecord Flag to control incoming message recording
     * @param chatMemoryOutRecord Flag to control outgoing message recording
     */
    public AgentChatMemoryAdvisor(ChatMemory chatMemory, String systemTextAdvise, boolean chatMemoryInRecord,boolean chatMemoryOutRecord) {
        super(chatMemory);
        this.systemTextAdvise = systemTextAdvise;
        this.chatMemoryInRecord = chatMemoryInRecord;
        this.chatMemoryOutRecord = chatMemoryOutRecord;
    }

    public AgentChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
                                  String systemTextAdvise) {
        this(chatMemory, defaultConversationId, chatHistoryWindowSize, systemTextAdvise,
                Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER, true,true); // Default value is true
    }

    public AgentChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
                                  String systemTextAdvise, int order) {
        this(chatMemory, defaultConversationId, chatHistoryWindowSize, systemTextAdvise, order, true,true); // Default value is true
    }

    public AgentChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
                                  String systemTextAdvise, int order, boolean chatMemoryInRecord, boolean chatMemoryOutRecord) {
        super(chatMemory, defaultConversationId, chatHistoryWindowSize, true, order);
        this.systemTextAdvise = systemTextAdvise;
        this.chatMemoryInRecord = chatMemoryInRecord;
        this.chatMemoryOutRecord = chatMemoryOutRecord;
    }

    /**
     * Builder pattern creator
     * @param chatMemory The chat memory instance to use
     * @return Builder instance
     */
    public static AgentChatMemoryAdvisor.Builder builder(ChatMemory chatMemory) {
        return new AgentChatMemoryAdvisor.Builder(chatMemory);
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

        //advisedRequest = agentRound(advisedRequest);

        advisedRequest = this.before(advisedRequest);

        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);

        this.observeAfter(advisedResponse);

        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

        //advisedRequest = agentRound(advisedRequest);

        Flux<AdvisedResponse> advisedResponses = this.doNextWithProtectFromBlockingBefore(advisedRequest, chain,
                this::before);

        return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, this::observeAfter);
    }

    /**
     * Determines the current conversation round and adds it to the request
     * @param advisedRequest The request to process
     * @return Updated request with round information
     */
    private AdvisedRequest agentRound(AdvisedRequest advisedRequest) {

        String conversationId = this.doGetConversationId(advisedRequest.adviseContext());
        int round = 0;
        List<Message> messages = this.getChatMemoryStore().get(conversationId, 1);
        if (CollectionUtils.isEmpty(messages)) {
            // First round of conversation
             round = 1;
        } else {
            // Get the round number from the message metadata
            round = (int) messages.get(0).getMetadata().get("round");
            round = round + 1;
        }
        // Add round information to the context
        advisedRequest.adviseContext().put("round",round);
        return advisedRequest;
    }

    /**
     * Processes the request before sending it to the LLM
     * @param request The advised request to process
     * @return The processed advised request
     */
    private AdvisedRequest before(AdvisedRequest request) {

        // 1. Advise system parameters.
        List<Message> memoryMessages = this.getChatMemoryStore()
                .get(this.doGetConversationId(request.adviseContext()),
                        this.doGetChatMemoryRetrieveSize(request.adviseContext()));

//        String memory = (memoryMessages != null) ? memoryMessages.stream()
//                .filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
//                .map(m -> m.getMessageType() + ":" + ((Content) m).getText())
//                .collect(Collectors.joining(System.lineSeparator())) : "";
//
        Map<String, Object> advisedSystemParams = new HashMap<>(request.systemParams());
        advisedSystemParams.putAll(request.advisorParams());

        // 2. Advise the system text.
        String advisedSystemText = request.systemText() + System.lineSeparator() + this.systemTextAdvise;
        // Replace placeholders in advisedSystemText
        PromptTemplate promptTemplate = new SystemPromptTemplate(advisedSystemText);
        advisedSystemText = promptTemplate.render(advisedSystemParams);

        // 3. Create a new request with the advised system text and parameters.
        Message systemMessage = new SystemMessage(advisedSystemText);
        List<Message> messages = new ArrayList<>();
        messages.add(systemMessage);
        boolean isNeedAssistantMessage = "true".equals(request.adviseContext().getOrDefault(AgentsConstants.AgentsAdvisorKey.ASSISTANT_MESSAGE_HISTORY_NEED,"true"));
        boolean isNeedUserMessage = "true".equals(request.adviseContext().getOrDefault(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HISTORY_NEED,"true"));
        for (Message message : memoryMessages) {
            if(message.getMessageType() == MessageType.USER && isNeedUserMessage){
                messages.add(message);
            } else if(message.getMessageType() == MessageType.ASSISTANT && isNeedAssistantMessage){
                messages.add(message);
            }
        }

        String userText = request.userText();
        String userLastMessage = (String)advisedSystemParams.getOrDefault(AgentsConstants.AgentsConversationEnvironmentKey.USER_LAST_MESSAGE, userText);
        Assert.notNull(userLastMessage, "userLastMessage cannot be null");
        // If the last record is the same as the current and is a user message, remove it
        if( messages.size() > 0 ) {
            Message lastMessage = messages.get(messages.size() - 1);
            if(lastMessage.getMessageType() == MessageType.USER ){
                if(userLastMessage != null && userLastMessage.equals(lastMessage.getText())) {
                    messages.remove(messages.size() - 1);
                }
            }
        }

        AdvisedRequest advisedRequest = AdvisedRequest.from(request)
                .systemText("")
                .messages(messages)
                .build();
        // 4. If a user message template exists, render it accordingly
        if (request.advisorParams().containsKey(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE)) {
            String userMessageTemplate = (String) request.advisorParams().get(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE);
            int index = 1;
            // Avoid consecutive duplicate messages, process in reverse
            String lastMessage = userLastMessage;
            List<String> userMessageHis = new ArrayList<>();
            StringBuffer userMessageBuffer = new StringBuffer("");
            for (int i = memoryMessages.size() - 1 ; i >= 0 ; i--) {
                Message message = memoryMessages.get(i);
                // Only keep messages with different content
                if(message.getMessageType() == MessageType.USER){
                    if(lastMessage != null && !lastMessage.equals(message.getText()) && !userLastMessage.equals(message.getText())) {
                        userMessageHis.add(message.getText());
                        lastMessage = message.getText();
                    }
                }
            }
            if(userMessageHis.size() > 0) {
                userMessageBuffer.append("The user previously raised these requests:\n");
                for (int i = userMessageHis.size() - 1 ; i >= 0 ; i--) {
                    String message = userMessageHis.get(i);
                    userMessageBuffer.append(index++).append(".").append(message).append("\n");
                }
            }

            PromptTemplate userPromptTemplate = new SystemPromptTemplate(userMessageTemplate);
            Map<String, Object> userPromptParams = new HashMap<>(advisedSystemParams);
            userPromptParams.put(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HIS, userMessageBuffer.toString());
            userPromptParams.put(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_CURRENT, userText);
            String userNewText = userPromptTemplate.render(userPromptParams);
            if(StringUtils.hasText(userNewText)) {
                advisedRequest = AdvisedRequest.from(request)
                        .userText(userNewText)
                        .systemText("")
                        .messages(messages)
                        .build();
            }
        }
        // If context recording is not needed, return
        if(!this.chatMemoryInRecord){
            return advisedRequest;
        }

        // 5. Add the new user input to the conversation memory.
        UserMessage userMessage = new UserMessage(userLastMessage, request.media());
        // Also put the corresponding agent information in the metaData of userMessage
        userMessage.getMetadata().put("agentId",doGetAgentId(request.adviseContext()));
        //userMessage.getMetadata().put("round",doGetAgentRound(request.adviseContext()));
        this.getChatMemoryStore().add(this.doGetConversationId(request.adviseContext()), userMessage);

        return advisedRequest;
    }

    /**
     * Processes the response after receiving it from the LLM
     * @param advisedResponse The advised response to process
     */
    private void observeAfter(AdvisedResponse advisedResponse) {
        // If context recording is not needed, return
        if(!this.chatMemoryOutRecord){
            return;
        }
        // If the content is returned directly by the tool, do not record the context
        if(StringUtils.hasText(ChatResponseUtils.getFirstToolName(advisedResponse.response()))){
            return;
        }
        List<Message> assistantMessages = advisedResponse.response()
                .getResults()
                .stream()
                .map(g -> (Message) g.getOutput())
                .toList();
        // Also put the corresponding agent information in the metaData of assistantMessage
        String agentId = doGetAgentId(advisedResponse.adviseContext());
        //int round = doGetAgentRound(advisedResponse.adviseContext());
        assistantMessages.forEach(message -> {
            message.getMetadata().put("agentId",agentId);
            //message.getMetadata().put("round",round);
        });
        this.getChatMemoryStore().add(this.doGetConversationId(advisedResponse.adviseContext()), assistantMessages);
    }

    /**
     * Gets the agent ID from the context
     * @param context The context map
     * @return The agent ID or empty string if not found
     */
    protected String doGetAgentId(Map<String, Object> context) {
        return context.containsKey(CHAT_MEMORY_AGENT_ID_KEY)
                ? context.get(CHAT_MEMORY_AGENT_ID_KEY).toString() : "";
    }

    /**
     * Gets the agent round from the context
     * @param context The context map
     * @return The round number or 0 if not found
     */
    protected int doGetAgentRound(Map<String, Object> context) {
        return context.containsKey(CHAT_MEMORY_AGENT_ROUND)
                ? (int) context.get(CHAT_MEMORY_AGENT_ROUND) : 0;
    }

    /**
     * Builder class for AgentChatMemoryAdvisor
     */
    public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<ChatMemory> {

        private String systemTextAdvise = DEFAULT_SYSTEM_TEXT_ADVISE;

        protected Builder(ChatMemory chatMemory) {
            super(chatMemory);
        }

        /**
         * Set the system text advice for the advisor
         * @param systemTextAdvise The system text advice
         * @return The builder instance
         */
        public AgentChatMemoryAdvisor.Builder systemTextAdvise(String systemTextAdvise) {
            this.systemTextAdvise = systemTextAdvise;
            return this;
        }

        /**
         * Build the AgentChatMemoryAdvisor instance
         * @return The AgentChatMemoryAdvisor instance
         */
        public AgentChatMemoryAdvisor build() {
            return new AgentChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
                    this.systemTextAdvise, this.order);
        }

    }

}